import os
import shutil
import random
import numpy as np

"""
准备性工作
将PASCAL_VOC_2007拆分成训练集与测试集
"""

# 数据集基路径
BASE_PATH = './data/VOCdevkit/VOC2007/'
# 目标位置
TRG_PATH = './data/VOC_{}/{}'
TRG_PATH_ANNO = TRG_PATH + '/Annotations/'
TRG_PATH_JPEG = TRG_PATH + '/JPEGImages/'

# anno后缀
suffix_xml = '.xml'
# image后缀
suffix_img = '.jpg'

# 指定训练集目录文件表
xxx_train_path = BASE_PATH + 'ImageSets/Main/{}_train.txt'
# 指定测试集目录文件表
xxx_val_path = BASE_PATH + 'ImageSets/Main/{}_val.txt'

# 指定anno目录,及图片目录
voc_annotation_dir = BASE_PATH + 'Annotations/'
voc_jpeg_dir = BASE_PATH + 'JPEGImages/'


# ======================XXX_train/val.txt文件读取====================================
def parse_train_val(path):
    """
    根据路径提取对应的train/val文件内容
    :param path:
    :return:返回标签为1的样本集
    """
    samples = []

    with open(path, 'r') as f:
        line = f.readlines()
        for line in line:
            # 遍历每一行
            li = line.strip().split(' ')
            # 找到其中标签为1的行
            if len(li) == 3 and int(li[-1]) == 1:
                samples.append(li[0])

    tt = path.split('/')[-1].split('_')
    print('============================================')
    print('{} {} sample num is {}'.format(tt[0], tt[-1][:-4], len(samples)))
    return np.array(samples)


# =======================voc_xxx目录形成===================================
def sample_train_val(samples, class_name):
    """

    其作用是为每个类别的数据创建一个目录,(名为voc_xxx,xxx代表分类名)
    并且根据给出样本中的图片名从源目标集将图片及对应的anno拷贝到当前训练集

    :param samples: train,val样本dict('train':list,'val':list)
    :param class_name: 分类名用于生成对应名称的文件夹与xxx.csv文件
    :return:
    生成的文件结构如下
    - voc_xxx
        -train
            - Annotations
            - JPEGImages
            - xxx.csv
        -val
            - Annotations
            - JPEGImages
            - xxx.csv
    """

    # 遍历train与val对其进行处理
    for key in samples.keys():
        # 检查文件夹是否存在
        check_dir(TRG_PATH_ANNO.format(class_name, key))
        check_dir(TRG_PATH_JPEG.format(class_name, key))
        # 获取图片名list
        xxx_sample = samples.get(key)
        # 遍历所有的文件名,并进行拼接处理
        for sample_name in xxx_sample:
            # 拼接源anno路径
            src_anno_path_one = voc_annotation_dir + sample_name + suffix_xml
            # 拼接目标anno路径
            trg_anno_path_one = TRG_PATH_ANNO.format(class_name, key) + sample_name + suffix_xml

            # 拼接源JPEGImg路径
            src_jpeg_path_one = voc_jpeg_dir + sample_name + suffix_img
            # 拼接目标JPEGImg路径
            trg_jpeg_path_one = TRG_PATH_JPEG.format(class_name, key) + sample_name + suffix_img

            # 拷贝
            shutil.copyfile(src_anno_path_one, trg_anno_path_one)
            shutil.copyfile(src_jpeg_path_one, trg_jpeg_path_one)

        # 生成csv文件
        csv_path = TRG_PATH.format(class_name, key + '/') + class_name + '.csv'
        np.savetxt(csv_path, np.array(xxx_sample), fmt='%s')


# ======================检查文件夹是否存在====================================

def check_dir(path):
    if not os.path.isdir(path):
        os.makedirs(path)


# ==========================================================
def generate_dataset(name):
    """
    从src目录中抽取对应分类的数据,并生成VOC_XXX目录
    :param name:
    :return:
    """
    sample_dict_train_val = {
        'train': parse_train_val(xxx_train_path.format(name)),
        'val': parse_train_val(xxx_val_path.format(name))
    }
    sample_train_val(sample_dict_train_val, name)


# ==========================================================

def get_classes():
    """
    得到所有的分类名称
    :return:
    """
    image_main_path = BASE_PATH + 'ImageSets/Main/'
    class_set = set()
    for file in os.listdir(image_main_path):
        if len(file.split('_')) > 1:
            class_set.add(file.split('_')[0])
    return list(class_set)

# ==========================================================


if __name__ == '__main__':
    """
    现在要做的事情是将ImageSets/Main/目录下XXX_train.txt和XXX_val.txt剔除标签列(第二列)
    选择其中标签值为1的列,然后拷贝到TRG_PATH路径的train和val文件夹下
    """
    class_names = get_classes()
    for name in class_names:
        generate_dataset(name)
